package com.starhub.application.rag.processor.domain;

import dev.langchain4j.data.embedding.Embedding;

import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.Future;
import java.util.stream.Collectors;
import lombok.extern.slf4j.Slf4j;

import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import com.starhub.application.collection.entity.KnowledgeCollection;
import com.starhub.application.collection.enums.MetaDataCollectionTypeEnums;
import com.starhub.application.collection.service.DataManipulationService;
import com.starhub.application.collection.service.KnowledgeCollectionService;
import com.starhub.application.rag.dto.MetaRagDto;
import com.starhub.application.rag.processor.EmbeddingProcessor;
import com.starhub.application.rag.result.RetrieveResult;
import com.starhub.application.rag.transformer.SimpleTextSegmentTransformer;
import com.starhub.application.rag.util.MilvusConditionBuilder;
import com.starhub.utils.Util;
import com.starhub.application.rag.config.ThreadPoolConfig;

/**
 * 向量化处理器
 * 用于处理向量化相关的业务逻辑
 * 
 * @version: V1.0
 * @author: liuhf
 * @time 2025-01-27
 */
@Component
@Slf4j
public class EmbeddingConditionProcessor extends EmbeddingProcessor {

    @Autowired
    SimpleTextSegmentTransformer simpleTextSegmentTransformer;

    @Autowired
    private KnowledgeCollectionService knowledgeCollectionService;

    @Autowired
    private DataManipulationService dataManipulationService;

    @Autowired
    private MilvusConditionBuilder milvusConditionBuilder;

    @Autowired
    private ThreadPoolConfig threadPoolConfig;

    /**
     * 向量化检索 DB检索模式 直接返回condition
     * 
     * @param metaRagDto
     * @param vectorModel
     * @return
     */
    public String embeddingRetrievalCondition(MetaRagDto metaRagDto, String vectorModel, Embedding questionEmbedding) {
        // 获取元数据表关联
        List<KnowledgeCollection> metaCollections = knowledgeCollectionService
                .selectBatchIds(metaRagDto.getMetaCollectionIds());

        // 执行每个元数据表的检索任务
        List<Future<RetrieveResult>> futures = metaCollections.stream()
                .map(collection -> threadPoolConfig.getThreadPoolExecutor().submit(
                        () -> {
                            return executeEmbeddingRetrieval(collection, metaRagDto, vectorModel,
                                    questionEmbedding);
                        }))
                .collect(Collectors.toList());

        List<RetrieveResult> retrieveResults = futures.stream().map(future -> {
            try {
                return future.get();
            } catch (Exception e) {
                log.error("检索任务执行失败", e);
            }
            return null;
        }).collect(Collectors.toList());
        return milvusConditionBuilder.buildNormalCondition(retrieveResults);
    }

    /**
     * 执行单个集合的检索任务
     * 
     * @param collection        知识库集合
     * @param metaRagDto        元数据检索请求
     * @param vectorModel       向量模型标识
     * @param questionEmbedding 问题向量
     * @return 检索结果
     */
    private RetrieveResult executeEmbeddingRetrieval(KnowledgeCollection collection,
            MetaRagDto metaRagDto,
            String vectorModel,
            Embedding questionEmbedding) {
        try {
            // 对应的元数据是向量存储模式
            if (collection.getMetaStorageType() == MetaDataCollectionTypeEnums.VECTOR_REPRESENTATION.getCode()) {
                // 获取检索结果
                List<String> itemIds = this.embeddingRetrieval(collection, metaRagDto, vectorModel, questionEmbedding);     

                // 获取有关联的主数据ID
                List<String> mainIds = new ArrayList<>();
                // 数据模式解析

                // 获取有关联的主数据ID 这里的检索目前直接根据关联表检索，等于查出所有有关的mainid，后续考虑如何优化
                mainIds = dataManipulationService.getMainDatas(metaRagDto.getMainCollectionId(),
                        Util.null2String(collection.getId()),
                        itemIds);

                return RetrieveResult.builder()
                        .collectionId(collection.getId())
                        .knowledgeCollection(collection)
                        .itemIds(itemIds)
                        .mainIds(mainIds)
                        .build();
            } else if (collection.getMetaStorageType() == MetaDataCollectionTypeEnums.DATABASE_REPRESENTATION
                    .getCode()) {
                // TODO 存储类型为数据表示
                return null;
            } else {
                return null;
            }

        } catch (Exception e) {
            log.error("集合 {} 检索失败: {}", collection.getCollectionIdentifier(), e.getMessage(), e);
            return null;
        }
    }
}
